# 从网上找到并魔改后能用的程序，可以在自己程序中使用与yolov5\detect.py相同的检测流程。完全可用
# 导入需要的库
import os
import sys
from pathlib import Path
import numpy as np
import cv2
import torch
import torch.backends.cudnn as cudnn

# 初始化目录
ROOT = "..\\YOLO\\yolov5"  # 若不在根目录下运行程序此处会有问题
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # 将YOLOv5的根目录添加到环境变量中（程序结束后删除）
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
 
from models.common import DetectMultiBackend
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr,
                           increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from utils.torch_utils import select_device, time_sync
from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox, mixup, random_perspective
 
weights = ROOT / "50l.pt"  # 权重文件地址

imgsz=(640, 480)  # 输入图片的大小 默认640(pixels)
conf_thres=0.25  # object置信度阈值 默认0.25  用在nms中
iou_thres=0.45  # 做nms的iou阈值 默认0.45   用在nms中
max_det=1000  # 每张图片最多的目标数量  用在nms中
device="0"  # 设置代码执行的设备 cuda device, i.e. 0 or 0,1,2,3 or cpu
classes=None  # 在nms中是否是只保留某些特定的类 默认是None 就是所有类只要满足条件都可以保留 --class 0, or --class 0 2 3
agnostic_nms=False  # 进行nms是否也除去不同类别之间的框 默认False
augment=False  # 预测是否也要采用数据增强 TTA 默认False
visualize=False  # 特征图可视化 默认FALSE
half=False  # 是否使用半精度 Float16 推理 可以缩短推理时间 但是默认是False
dnn=False  # 使用OpenCV DNN进行ONNX推理
 
# 获取设备
device = select_device(device)
 
def LoadModel():
    global model, stride, names, pt, imgsz, half
    # 载入模型
    model = DetectMultiBackend(weights, device=device, dnn=dnn)

    stride, names, pt, jit, onnx, engine = model.stride, model.names, model.pt, model.jit, model.onnx, model.engine
    imgsz = check_img_size(imgsz, s=stride)  # 检查图片尺寸
    
    # 使用半精度 Float16 推理
    half &= (pt or jit or onnx or engine) and device.type != "cpu"  # FP16 supported on limited backends with CUDA
    if pt or jit:
        model.model.half() if half else model.model.float()
    return model

def detect(img):
    # 开始预测
    if half != False:
        model.warmup(imgsz=(1, 3, *imgsz), half=half)  # warmup
    else:
        model.warmup(imgsz=(1, 3, *imgsz))  # warmup
    #对图片进行处理
    im0 = img
    # Padded resize
    im = letterbox(im0, imgsz, stride, auto=pt)[0]
    # Convert
    im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
    im = np.ascontiguousarray(im)
    im = torch.from_numpy(im).to(device)
    im = im.half() if half else im.float()  # uint8 to fp16/32
    im /= 255  # 0 - 255 to 0.0 - 1.0
    if len(im.shape) == 3:
        im = im[None]  # expand for batch dim
    # 预测
    pred = model(im, augment=augment, visualize=visualize)
    # NMS
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
    #用于存放结果
    detections=[]
    # Process predictions
    for i, det in enumerate(pred):  # per image 每张图片
        if len(det):
            # Rescale boxes from img_size to im0 size
            det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
            # 写入结果
            for *xyxy, conf, cls in reversed(det):
                detections.append([float(xyxy[0]), float(xyxy[1]), float(xyxy[2]), float(xyxy[3]), float(conf), int(cls)])
    return detections

if __name__ == "__main__":
    LoadModel()
    img = cv2.imread("test2.jpg")
    detect(img)